# 基于cnn实现猫狗识别
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

# 1、加载数据
root_dir = r"/home/hy/vscode/pytorch_by_example/data"
img_dir = os.path.join(root_dir, "imgs")
path_txt_train = os.path.join(root_dir, "labels", "train.txt")
path_txt_valid = os.path.join(root_dir, "labels", "valid.txt")
transforms_func = transforms.Compose([
    transforms.Resize((8, 8)),
    transforms.ToTensor(),
])
class COVID19Dataset(Dataset):
        def __init__(self, root_dir, txt_path, transform=None):
            """
            获取数据集的路径、预处理的方法
            """
            self.root_dir = root_dir
            self.txt_path = txt_path
            self.transform = transform
            self.img_info = []  # [(path, label), ... , ]
            self.label_array = None
            self._get_img_info()

        def __getitem__(self, index):
            """
            输入标量index, 从硬盘中读取数据，并预处理，to Tensor
            :param index:
            :return:
            """
            path_img, label = self.img_info[index]
            img = Image.open(path_img).convert('L')

            if self.transform is not None:
                img = self.transform(img)

            return img, label

        def __len__(self):
            if len(self.img_info) == 0:
                raise Exception("\ndata_dir:{} is a empty dir! Please checkout your path to images!".format(
                    self.root_dir))  # 代码具有友好的提示功能，便于debug
            return len(self.img_info)

        def _get_img_info(self):
            """
            实现数据集的读取，将硬盘中的数据路径，标签读取进来，存在一个list中
            path, label
            :return:
            """
            # 读取txt，解析txt
            with open(self.txt_path, "r") as f:
                txt_data = f.read().strip()
                txt_data = txt_data.split("\n")

            self.img_info = [(os.path.join(self.root_dir, i.split()[0]), int(i.split()[2])) for i in txt_data]
    
# DataLoader批处理
train_data = COVID19Dataset(root_dir=img_dir, txt_path=path_txt_train, transform=transforms_func)
valid_data = COVID19Dataset(root_dir=img_dir, txt_path=path_txt_valid, transform=transforms_func)
train_loader = DataLoader(dataset=train_data, batch_size=2)
valid_loader = DataLoader(dataset=valid_data, batch_size=2)

# 2、 定义CNN模型模块
class TinnyCNN(nn.Module):
    def __init__(self):
        super(TinnyCNN, self).__init__()
        # 定义卷积层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)  # 输入1通道，输出32通道
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # 输入32通道，输出64通道
        # 定义全连接层
        self.fc1 = nn.Linear(64 * 2 * 2, 128)  # 展平后输入到全连接层
        self.fc2 = nn.Linear(128, 2)  # 2 个类别

    def forward(self, x):
        x = F.relu(self.conv1(x))  # 第一层卷积 + ReLU
        x = F.max_pool2d(x, 2)     # 最大池化
        x = F.relu(self.conv2(x))  # 第二层卷积 + ReLU
        x = F.max_pool2d(x, 2)     # 最大池化
        x = x.view(x.size(0), -1)  # 动态展平
        x = F.relu(self.fc1(x))    # 全连接层 + ReLU
        x = self.fc2(x)            # 最后一层输出
        return x

model = TinnyCNN()


# 3、定义损失函数和优化器
cro_loss = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 4、模型训练
num_epochs = 100
model.train()
for epoch in range(num_epochs):
    total_loss = 0
    for data, labels in train_loader:
        #前向传播
        outputs = model(data)
        loss = cro_loss(outputs, labels)

        #反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step() # 更新参数

        total_loss += loss.item()
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(train_loader):.4f}")



# 5、模型测试
model.eval()
correct = 0
total = 0
with torch.no_grad(): #关闭梯度计算
    for data, label in valid_loader:
        outputs = model(data)
        # 计算分类准确率
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")       

# 6. 可视化测试结果
dataiter = iter(valid_loader)
images, labels = next(dataiter)
outputs = model(images)
_, predictions = torch.max(outputs, 1)

fig, axes = plt.subplots(1, 6, figsize=(12, 4))
for i in range(2):
    axes[i].imshow(images.data[i][0], cmap='gray')
    axes[i].set_title(f"Label: {labels[i]}\nPred: {predictions[i]}")
    axes[i].axis('off')
plt.show()